from keras.datasets.mnist import load_data
from keras import Sequential, models, Model, layers, activations, optimizers, losses
(x_train, y_train), (x_test, y_test) = load_data()
x_train = x_train.reshape(-1, 28, 28, 1) / 255
x_test = x_test.reshape(-1, 28, 28, 1) / 255

class VGG16(Model):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conv = Sequential([
            layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=64, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=64, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=64, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=128, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=128, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.Conv2D(filters=128, kernel_size=(3, 3), padding='same', activation=activations.relu),
            layers.MaxPooling2D()
        ])
        self.flat = Sequential(layers.Flatten())
        self.fc = Sequential([
            layers.Dense(units=4096, activation=activations.relu),
            layers.Dense(units=4096, activation=activations.relu),
            layers.Dense(units=10, activation=activations.softmax)
        ])

    def call(self, inputs, training=None, mask=None):
        out = self.conv(inputs)
        out = self.flat(out)
        out = self.fc(out)
        return out


model = VGG16()
model.build(input_shape=[None, 28, 28, 1])
model.summary()

model.compile(optimizer=optimizers.Adam(), loss=losses.sparse_categorical_crossentropy, metrics='acc')
model.fit(x_train, y_train, batch_size=100, epochs=1)
model.evaluate(x_test, y_test)
